"""
核反应堆子通道分析程序 - 主程序入口
"""

import os
import sys
import argparse
from core.analysis import SubchannelAnalysis
from io.file_operations import (
    create_results_folder,
    save_history_data,
    export_results_summary
)
from plotting.plot_functions import (
    plot_results,
    plot_transient_results,
    plot_channel_type_comparison,
    plot_convergence_history,
    plot_heat_transfer_analysis
)

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='核反应堆子通道分析程序')
    
    # 添加命令行参数
    parser.add_argument('--n_axial', type=int, default=2000,
                      help='轴向节点数 (默认: 2000)')
    parser.add_argument('--total_time', type=float, default=5.0,
                      help='总计算时间 (s) (默认: 5.0)')
    parser.add_argument('--time_step', type=float, default=1e-4,
                      help='时间步长 (s) (默认: 1e-4)')
    parser.add_argument('--save_interval', type=int, default=100,
                      help='保存间隔 (默认: 100)')
    parser.add_argument('--results_folder', type=str, default='results',
                      help='结果保存文件夹 (默认: results)')
    parser.add_argument('--plot_only', action='store_true',
                      help='仅绘制已有结果')
    parser.add_argument('--resume', type=str,
                      help='从指定的状态文件恢复计算')
    
    return parser.parse_args()

def main():
    """主程序入口"""
    # 解析命令行参数
    args = parse_arguments()
    
    try:
        # 创建结果文件夹
        create_results_folder(args.results_folder)
        
        if not args.plot_only:
            # 创建分析实例
            analysis = SubchannelAnalysis(
                n_axial=args.n_axial,
                total_time=args.total_time,
                time_step=args.time_step,
                save_interval=args.save_interval,
                results_folder=args.results_folder
            )
            
            # 如果指定了恢复文件，从文件恢复状态
            if args.resume:
                analysis.load_state(args.resume)
            
            # 运行瞬态模拟
            print("开始瞬态计算...")
            analysis.run_transient_simulation()
            print("瞬态计算完成")
            
            # 保存历史数据
            save_history_data(analysis)
            
            # 导出结果摘要
            export_results_summary(analysis, 
                                 os.path.join(args.results_folder, 'summary.txt'))
            
        else:
            # 从最新的状态文件加载结果
            state_files = [f for f in os.listdir(args.results_folder) 
                         if f.startswith('full_state_') and f.endswith('.npz')]
            if not state_files:
                print("错误：未找到状态文件")
                return
                
            latest_state = sorted(state_files)[-1]
            analysis = SubchannelAnalysis()
            analysis.load_state(os.path.join(args.results_folder, latest_state))
        
        # 绘制结果
        print("\n绘制计算结果...")
        
        # 1. 绘制基本结果
        plot_results(analysis)
        
        # 2. 绘制瞬态结果
        plot_transient_results(analysis)
        
        # 3. 绘制子通道类型对比
        plot_channel_type_comparison(analysis)
        
        # 4. 绘制收敛历史
        plot_convergence_history(analysis)
        
        # 5. 绘制传热分析
        plot_heat_transfer_analysis(analysis)
        
        print("所有结果已保存到:", args.results_folder)
        
    except Exception as e:
        print(f"程序运行出错: {str(e)}")
        import traceback
        print(traceback.format_exc())
        return 1
        
    return 0

if __name__ == "__main__":
    sys.exit(main())
